from ultralytics import YOLO
from ultralytics.engine.results import Results
import csv


def predict_and_save(source_filename, model_name, out_name):
    model = YOLO(model_name)
    fpout = open(out_name, 'w')
    writer = csv.DictWriter(fpout, fieldnames=[
        'frame', 'obj', 'label', 'x', 'y', 'confident',
    ])
    writer.writeheader()

    results = model(source_filename, stream=True, save=True)
    
    for iframe, res in enumerate(results):
        res:Results
        kp = res.keypoints
        kp = kp.cpu().numpy()

        if kp.conf is None:
            print(f'No result for {iframe=}')
            continue
        
        for iobj in range(kp.shape[0]):
            for ipos in range(kp.shape[1]):
                xy = kp.xy[iobj, ipos, :]
                writer.writerow({
                    'frame': iframe,
                    'obj': iobj,
                    'label': ipos,
                    'x': xy[0],
                    'y': xy[1],
                    'confident': kp.conf[iobj, ipos]
                })


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--model-name', '-m', default='yolo11m-pose.pt')
    parser.add_argument('--source', '-s', required=True)
    parser.add_argument('--output', '-o', default='keypoints.csv')

    p = parser.parse_args()
    predict_and_save(p.source, p.model_name, p.output)


